import traceback

import numpy as np
from sklearn.preprocessing import LabelEncoder, OrdinalEncoder
from tqdm import tqdm

from src.bean.vector_name_enum import VectorNameEnum
from src.model.model_common_util import load_vectors, save_vector, drop_and_get_cols, fill_bool_with_num
from src.model.v2.pre_handle_str import fill_str_with_blank
from src.util.common_util import printx


def handle_array(obj):
    v = 'blank'
    if obj is not np.nan:
        try:
            v = "|".join(f"{obj}"
                         .replace("]", "")
                         .replace("[", "")
                         .replace("\\", "")
                         .replace("\"", "")
                         .replace(" ", "")
                         .replace("\'", "")
                         .split(","))
        except Exception:
            printx(f"{obj}")
            traceback.print_exc()
    return v


def handle_series(series):
    printx(series.name)
    return series.progress_apply(handle_array)


def v1_pre_handle(df, model_id, encoder_type: VectorNameEnum = VectorNameEnum.v1_ordinal):
    try:
        df, num_cols, str_cols, bool_cols, array_cols = drop_and_get_cols(df)
        if num_cols:
            df.loc[:, num_cols] = df[num_cols].fillna(-1).astype(int)
        df = fill_str_with_blank(df, str_cols)
        fill_bool_with_num(df, bool_cols)
        if array_cols:
            tqdm.pandas()
            df.loc[:, array_cols] = df[array_cols].apply(handle_series).fillna('blank').astype(str)
        enc_cols = str_cols + array_cols
        vectors = load_vectors(encoder_type.value, model_id)
        if vectors is None:
            vectors = {}
        if encoder_type.value == 'label':
            for column in enc_cols:
                if column in vectors:
                    df.loc[:, column] = vectors[column].transform(df[column])
                else:
                    encoder = LabelEncoder()
                    df.loc[:, column] = encoder.fit_transform(df[column])
                    vectors[column] = encoder
        elif encoder_type.value == 'ordinal':
            if "encoder" in vectors:
                encoder = vectors["encoder"]
                enc_cols_sorted = [col for col in encoder.feature_names_in_ if col in df.columns]
                df.loc[:, enc_cols_sorted] = encoder.transform(df[enc_cols_sorted])
            else:
                encoder = OrdinalEncoder(handle_unknown="use_encoded_value", unknown_value=-1)
                df.loc[:, enc_cols] = encoder.fit_transform(df[enc_cols])
                vectors["encoder"] = encoder
        elif encoder_type.value == 'hdbscan':
            pass
        save_vector(vectors, encoder_type.value, model_id)
    except Exception as e:
        traceback.print_exc()
        printx(f"预处理报错:{e}")
